#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@File    : momentum.py
@Time    : 2024/12/01
@Author  : Liuli
@Desc    : 技术分析指标计算
"""

import numpy as np
import pandas as pd


def calculate_ma_crossover_signals(price_data, short_period=60, long_periods=[120, 180, 240, 300, 360, 420]):
    """
    计算均线死叉信号
    当短期均线(60日)连续死叉多个长期均线时，提示趋势破位
    
    Parameters:
    -----------
    price_data : pd.DataFrame
        资产价格数据，index为日期
    short_period : int
        短期均线周期（日）
    long_periods : list
        长期均线周期列表（日）
        
    Returns:
    --------
    pd.DataFrame
        死叉信号，1表示死叉，0表示非死叉
    """
    # 计算短期均线
    short_ma = price_data.rolling(window=short_period).mean()
    
    # 计算各个长期均线的死叉信号
    signals = pd.DataFrame(index=price_data.index, columns=price_data.columns)
    
    for col in price_data.columns:
        # 计算所有长期均线
        long_ma_signals = []
        for period in long_periods:
            long_ma = price_data[col].rolling(window=period).mean()
            # 当短期均线位于长期均线下方时，记为死叉
            death_cross = (short_ma[col] < long_ma) & (short_ma[col].shift(1) >= long_ma.shift(1))
            long_ma_signals.append(death_cross)
        
        # 只要有一个长期均线发生死叉，就记为信号
        signals[col] = pd.concat(long_ma_signals, axis=1).any(axis=1).astype(int)
    
    return signals


def calculate_macd_signals(price_data, fast=12, slow=26, signal=9):
    """
    计算MACD信号
    当MACD处于绿柱且绿柱扩大时，提示趋势破位
    
    Parameters:
    -----------
    price_data : pd.DataFrame
        资产价格数据，index为日期
    fast : int
        快线周期（日）
    slow : int
        慢线周期（日）
    signal : int
        信号线周期（日）
        
    Returns:
    --------
    pd.DataFrame
        MACD信号，1表示趋势破位，0表示非破位
    """
    signals = pd.DataFrame(index=price_data.index, columns=price_data.columns)
    
    for col in price_data.columns:
        # 计算EMA
        exp1 = price_data[col].ewm(span=fast, adjust=False).mean()
        exp2 = price_data[col].ewm(span=slow, adjust=False).mean()
        
        # 计算MACD线和信号线
        macd = exp1 - exp2
        signal_line = macd.ewm(span=signal, adjust=False).mean()
        
        # 计算MACD柱
        histogram = macd - signal_line
        
        # 当MACD柱为负且比前一日更负时，视为趋势破位信号
        signals[col] = ((histogram < 0) & (histogram < histogram.shift(1))).astype(int)
    
    return signals


def calculate_bollinger_signals(price_data, window=60, num_std=2):
    """
    计算布林带信号
    当价格从上到下穿越中轨时提示趋势破位，但突破下轨时不产生信号
    
    Parameters:
    -----------
    price_data : pd.DataFrame
        资产价格数据，index为日期
    window : int
        移动平均窗口期（日）
    num_std : int
        标准差倍数
        
    Returns:
    --------
    pd.DataFrame
        布林带信号，1表示趋势破位，0表示非破位
    """
    signals = pd.DataFrame(index=price_data.index, columns=price_data.columns)
    
    for col in price_data.columns:
        # 计算移动平均线（中轨）
        middle_band = price_data[col].rolling(window=window).mean()
        
        # 计算标准差
        std = price_data[col].rolling(window=window).std()
        
        # 计算上下轨
        upper_band = middle_band + num_std * std
        lower_band = middle_band - num_std * std
        
        # 计算信号：
        # 1. 从上到下穿越中轨时产生信号
        cross_middle = (price_data[col] < middle_band) & (price_data[col].shift(1) >= middle_band)
        # 2. 但是如果突破下轨，则不产生信号
        below_lower = price_data[col] < lower_band
        
        signals[col] = (cross_middle & ~below_lower).astype(int)
    
    return signals


def calculate_technical_signals(price_data):
    """
    综合计算技术指标信号
    只要有一个指标提示趋势破位，就认为该资产存在技术面风险
    
    Parameters:
    -----------
    price_data : pd.DataFrame
        资产价格数据，index为日期
        
    Returns:
    --------
    pd.DataFrame
        综合技术信号，1表示存在风险，0表示无风险
    """
    # 计算各个技术指标的信号
    ma_signals = calculate_ma_crossover_signals(price_data)
    macd_signals = calculate_macd_signals(price_data)
    bollinger_signals = calculate_bollinger_signals(price_data)
    
    # 综合三个信号：任意一个指标产生信号就视为趋势破位
    technical_signals = (ma_signals | macd_signals | bollinger_signals).astype(int)
    
    return technical_signals


if __name__ == "__main__":
    # 测试代码
    from WindPy import w
    w.start()
    
    # 获取测试数据
    start = '2020-01-01'
    end = '2024-01-19'
    index = ['H00300.CSI', 'H00905.CSI', 'H00852.SH', 'CBA06501.CS', 'CBA04601.CS',
             'CBA14101.CS', 'CBA20901.CS', 'SP500TR.SPI', 'AU.SHF', 'CU.SHF', 'SC.INE']
    
    data = w.wsd(index, "close", start, end, "")
    price_data = pd.DataFrame(np.array(data.Data).T, index=data.Times, columns=data.Codes)
    
    # 计算技术信号
    technical_signals = calculate_technical_signals(price_data)
    
    # 输出结果
    print("\n技术信号汇总：")
    print(technical_signals.tail())
    
    # 分析各个指标的信号触发频率
    ma_signals = calculate_ma_crossover_signals(price_data)
    macd_signals = calculate_macd_signals(price_data)
    bollinger_signals = calculate_bollinger_signals(price_data)
    
    print("\n各指标信号触发频率：")
    print("均线系统：", ma_signals.mean().mean())
    print("MACD系统：", macd_signals.mean().mean())
    print("布林带系统：", bollinger_signals.mean().mean())